{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cf149e0f-0c8d-486b-ba02-7c6982a6c463",
   "metadata": {},
   "source": [
    "# Finetuning E5-Large/LLaMA-3.2-1B Model on AllNLI for Embedding Tasks  \n",
    "\n",
    "## Goal  \n",
    "While the `intfloat/e5-large-v2 or LLaMA-3.2-1B` model is a powerful pretrained embedding model released on Hugging Face, adapting it for **specific downstream tasks** (such as semantic search, document retrieval, clustering, or retrieval-augmented generation) greatly benefits from **domain-specific finetuning**.  \n",
    "\n",
    "In this tutorial, we will demonstrate how to:  \n",
    "- Convert the Hugging Face `e5-large-v2/ LLaMA-3.2-1B` model into NeMo’s `.nemo` format.  \n",
    "- Prepare the **AllNLI triplet dataset** in a format compatible with NeMo’s `CustomRetrievalDataModule`.  \n",
    "- Fine-tune the E5 model to enhance its performance on **embedding-rich tasks** (retrieval, RAG, text similarity, etc.).  \n",
    "\n",
    "By leveraging **triplet training** (query, positive doc, negative doc), the model learns:  \n",
    "- To generate **semantically meaningful dense vector representations (embeddings)**.  \n",
    "- Improve retrieval quality by maximizing similarity between queries and positive docs while minimizing similarity with negatives.  \n",
    "\n",
    "***\n",
    "\n",
    "## NeMo Tools and Resources  \n",
    "- [NeMo Framework Documentation](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)  \n",
    "- Hugging Face Model Hub (`intfloat/e5-large-v2 OR meta-llama/Llama-3.2-1B`)  \n",
    "- NeMo `llm.import_ckpt` utility for checkpoint conversion  \n",
    "- NeMo `CustomRetrievalDataModule` for embedding retraining  \n",
    "\n",
    "***\n",
    "\n",
    "## Software Requirements  \n",
    "- NVIDIA NeMo Framework (`pip install nemo_toolkit[all]` or use NGC container)  \n",
    "- Hugging Face CLI & `datasets` library  \n",
    "- PyTorch >= 2.0, CUDA-enabled environment  \n",
    "\n",
    "***\n",
    "\n",
    "## Hardware Requirements  \n",
    "This playbook has been tested on:  \n",
    "- **Single GPU** setups (A100, H100) for quick runs  \n",
    "- **Multi-GPU** (2×A100, 2×H100) for large-scale training  \n",
    "- NeMo is fully scalable across **multi-node GPU clusters** via `torchrun` or `nemo_run`  \n",
    "\n",
    "***\n",
    "\n",
    "## Launching the NeMo Container  \n",
    "\n",
    "### With Docker  \n",
    "```bash\n",
    "docker run \\\n",
    "  --gpus all \\\n",
    "  --shm-size=2g \\\n",
    "  --net=host \\\n",
    "  --ulimit memlock=-1 \\\n",
    "  --rm -it \\\n",
    "  -v ${PWD}:/workspace \\\n",
    "  -w /workspace \\\n",
    "  nvcr.io/nvidia/nemo:25.04 bash\n",
    "```\n",
    "\n",
    "Start Jupyter inside the container:  \n",
    "```bash\n",
    "jupyter-lab --ip=0.0.0.0 --allow-root \\\n",
    "  --NotebookApp.token=\"embedding\" \\\n",
    "  --port=9989 --notebook-dir=/workspace\n",
    "```\n",
    "\n",
    "### With Enroot (alternative)  \n",
    "```bash\n",
    "mkdir -p \"$PWD/.jupyter_data\" \"$PWD/.jupyter_runtime\" \"$PWD/.jupyter_config\" \\\n",
    "         \"$PWD/.hf_cache\" \"$PWD/.matplotlib\" \"$PWD/.triton\" \"$PWD/.cache\" \\\n",
    "         \"$PWD/enroot_data\" \"$PWD/enroot_cache\" \"$PWD/nemo_home\" \\\n",
    "         \"$PWD/nemo_cache\" \"$PWD/nemo_run\"\n",
    "\n",
    "ENROOT_DATA_PATH=\"$PWD/enroot_data\" \\\n",
    "ENROOT_CACHE_PATH=\"$PWD/enroot_cache\" \\\n",
    "enroot start --root \\\n",
    "  --mount \"$PWD:/host_pwd\" \\\n",
    "  --env NVIDIA_VISIBLE_DEVICES=0,1 \\\n",
    "  --env NVIDIA_DRIVER_CAPABILITIES=all \\\n",
    "  --env JUPYTER_DATA_DIR=/host_pwd/.jupyter_data \\\n",
    "  --env JUPYTER_RUNTIME_DIR=/host_pwd/.jupyter_runtime \\\n",
    "  --env JUPYTER_CONFIG_DIR=/host_pwd/.jupyter_config \\\n",
    "  --env HF_HOME=/host_pwd/.hf_cache \\\n",
    "  --env MPLCONFIGDIR=/host_pwd/.matplotlib \\\n",
    "  --env TRITON_CACHE_DIR=/host_pwd/.triton \\\n",
    "  --env XDG_CACHE_HOME=/host_pwd/.cache \\\n",
    "  --env NEMO_HOME=/host_pwd/PythonNotebook/nemo_home \\\n",
    "  --env NEMO_MODELS_CACHE=/host_pwd/PythonNotebook/nemo_cache \\\n",
    "  --env NEMO_RUN_DIR=/host_pwd/PythonNotebook/nemo_run \\\n",
    "  nemo-25.04 \\\n",
    "  jupyter-lab --ip=0.0.0.0 --allow-root \\\n",
    "              --NotebookApp.token=\"embedding\" \\\n",
    "              --port=9989 \\\n",
    "              --notebook-dir=/host_pwd\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fee26a79-f000-41be-98b7-2c72653cc4cb",
   "metadata": {},
   "source": [
    "# Prepare AllNLI triplet data for CustomRetrievalDataModule\n",
    "\n",
    "This cell:\n",
    "- Downloads the AllNLI triplet training split via Hugging Face Datasets.\n",
    "- Transforms each triplet into a record compatible with CustomRetrievalDataModule:\n",
    "  - query: anchor sentence\n",
    "  - pos_doc: positive (entailing) sentence\n",
    "  - neg_doc: negative (contradicting/neutral) sentence\n",
    "- Saves all records to a UTF-8 JSON file allnli_triplet.json (pretty-printed).\n",
    "\n",
    "Prerequisites:\n",
    "- pip install datasets\n",
    "- Internet access and sufficient disk space (dataset is large).\n",
    "\n",
    "Output format (example record):\n",
    "```\n",
    "{\n",
    "  \"query\": \"A man is playing a guitar on stage.\",\n",
    "  \"pos_doc\": \"Someone is performing music in front of an audience.\",\n",
    "  \"neg_doc\": \"No one is playing any instruments.\"\n",
    "}\n",
    "```\n",
    "\n",
    "Notes:\n",
    "- To test quickly, you can load a subset: split='train[:1%]'.\n",
    "- If you hit memory limits, consider processing in chunks or using streaming (IterableDataset) and writing incrementally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9ecd43e-20cf-4b1c-bbe6-fd4cd41305a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from datasets import load_dataset\n",
    "\n",
    "print(\"Downloading AllNLI triplet dataset (train split)...\")\n",
    "ds = load_dataset('sentence-transformers/all-nli', 'triplet', split='train')\n",
    "len(ds), ds\n",
    "\n",
    "\n",
    "print(\"Transforming to CustomRetrievalDataModule-compatible JSON...\")\n",
    "\n",
    "records = [\n",
    "    {\n",
    "        \"query\": ex[\"anchor\"],\n",
    "        \"pos_doc\": ex[\"positive\"],\n",
    "        \"neg_doc\": ex[\"negative\"],\n",
    "    }\n",
    "    for ex in ds\n",
    "]\n",
    "\n",
    "out_path = \"allnli_triplet.json\"\n",
    "with open(out_path, \"w\", encoding=\"utf-8\") as f:\n",
    "    json.dump(records, f, ensure_ascii=False, indent=2)\n",
    "\n",
    "print(f\"Saved {len(records)} triplets to {out_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78407b5d-4867-4065-abcd-f64683203b12",
   "metadata": {},
   "source": [
    "# Import Hugging Face E5-Large model into NeMo format\n",
    "\n",
    "This script:\n",
    "- Uses NVIDIA NeMo's `llm.import_ckpt` utility to **download and convert** the Hugging Face model `intfloat/e5-large-v2`.\n",
    "- Wraps the E5-Large embedding model in NeMo’s `BertEmbeddingModel` + `BertEmbeddingLargeConfig`.\n",
    "- Saves the converted checkpoint to a local `.nemo` file (`e5-large-v2.nemo`).\n",
    "\n",
    "Steps performed:\n",
    "1. Define working directory as default `NEMO_HOME` and `NEMO_MODELS_CACHE`.\n",
    "2. Initialize the NeMo model configuration for E5-Large.\n",
    "3. Import the model from Hugging Face (`hf://intfloat/e5-large-v2`).\n",
    "4. Persist the converted checkpoint locally as `e5-large-v2.nemo`.\n",
    "5. Log start/end points for progress visibility.\n",
    "\n",
    "Prerequisites:\n",
    "- Nemo Container OR`pip install nemo_toolkit[all]` (or latest NeMo nightly with `collections.llm`).\n",
    "- Hugging Face access (ensure `huggingface_hub` is installed).\n",
    "- Adequate GPU/CPU + disk space (~1.5GB checkpoint).\n",
    "\n",
    "Usage:\n",
    "```bash\n",
    "python import_e5_large.py\n",
    "```\n",
    "\n",
    "Output:\n",
    "- A NeMo-compatible checkpoint: `e5-large-v2.nemo`  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "64fe4215-cc65-46ad-85a6-15e357d82c6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing import_e5_large.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile import_e5_large.py\n",
    "#!/usr/bin/env python3\n",
    "\"\"\"\n",
    "import_e5_large.py\n",
    "\n",
    "Downloads and converts the `intfloat/e5-large-v2` embedding model from Hugging Face\n",
    "into NeMo `.nemo` format using NeMo's llm.import_ckpt utility.\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "import logging\n",
    "from nemo.collections import llm\n",
    "\n",
    "# Setup logging\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "\n",
    "def main():\n",
    "    # Step 1: Define working directory and environment paths\n",
    "    cwd = os.getcwd()\n",
    "    os.environ.setdefault(\"NEMO_HOME\", cwd)\n",
    "    os.environ[\"NEMO_MODELS_CACHE\"] = cwd\n",
    "\n",
    "    # Step 2: Create model config for E5-Large embeddings\n",
    "    model_config = llm.BertEmbeddingModel(llm.BertEmbeddingLargeConfig())\n",
    "\n",
    "    # Hugging Face source\n",
    "    hf_source = 'hf://intfloat/e5-large-v2'\n",
    "\n",
    "    # Step 3: Convert and save model\n",
    "    output_file = os.path.join(cwd, 'e5-large-v2.nemo')\n",
    "    logging.info(f\" Importing from {hf_source} → {output_file}...\")\n",
    "\n",
    "    llm.import_ckpt(\n",
    "        model=model_config,\n",
    "        source=hf_source,\n",
    "        output_path=output_file,\n",
    "    )\n",
    "\n",
    "    # Step 4: Confirm success\n",
    "    logging.info(f\" Done. Checkpoint saved to {os.path.abspath(output_file)}\")\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3669f058-fe8e-4f78-9112-2c7a97ee1db8",
   "metadata": {},
   "source": [
    "# Authenticate with Hugging Face and run NeMo model import  \n",
    "\n",
    "This cell:  \n",
    "- Logs into Hugging Face with a personal access token so that private or gated models can be downloaded.  \n",
    "- Executes the `import_e5_large.py` script using `torchrun` — this ensures proper distributed/parallel execution if GPUs are available.  \n",
    "\n",
    "Steps:  \n",
    "1. `huggingface-cli login --token \"…\" `\n",
    "   - Provides your token to Hugging Face Hub (so `intfloat/e5-large-v2` can be accessed).  \n",
    "   - Token is stored locally (usually under `~/.huggingface/token`).  \n",
    "\n",
    "2. `torchrun import_e5_large.py`  \n",
    "   - Launches the import script you created.  \n",
    "   - Downloads the Hugging Face model weights.  \n",
    "   - Converts them into NeMo `.nemo` format.  \n",
    "   - Saves locally as `e5-large-v2.nemo`.  \n",
    "\n",
    "\n",
    "- Use `!huggingface-cli login` without `--token` and paste interactively, or  \n",
    "- Set the token in an environment variable (e.g., `export HUGGINGFACE_TOKEN=...`) and run `!huggingface-cli login --token \"$HUGGINGFACE_TOKEN\"`.  \n",
    "\n",
    "Output after successful run:  \n",
    "- `e5-large-v2.nemo` checkpoint will appear in your working directory, ready for use inside NeMo pipelines.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2de82074-aeb5-4e84-94d3-6e1508dd9124",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!huggingface-cli login --token \"hf_***************************\"\n",
    "!torchrun import_e5_large.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "20bfb0da-e950-436e-a912-7d49e8b5bbde",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing import_llama1b.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile import_llama1b.py\n",
    "#!/usr/bin/env python3\n",
    "\"\"\"\n",
    "import_llama1b.py\n",
    "\n",
    "This script downloads and converts a Hugging Face-hosted LLaMA-3 1B embedding model\n",
    "into NeMo format using NVIDIA NeMo's `llm.import_ckpt` utility.\n",
    "\n",
    "The final `.nemo` checkpoint is saved in the current working directory.\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "import logging\n",
    "from nemo.collections import llm\n",
    "\n",
    "# Setup logging\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "\n",
    "def main():\n",
    "    # Step 1: Define working directory and set required environment variables\n",
    "    cwd = os.getcwd()\n",
    "    os.environ.setdefault(\"NEMO_HOME\", cwd)\n",
    "    os.environ[\"NEMO_MODELS_CACHE\"] = cwd\n",
    "\n",
    "    # Step 2: Create model config for LLaMA-3 1B embeddings\n",
    "    model_config = llm.LlamaEmbeddingModel(llm.Llama32EmbeddingConfig1B())\n",
    "\n",
    "    # Define Hugging Face source\n",
    "    hf_source = 'hf://meta-llama/Llama-3.2-1B'\n",
    "\n",
    "    # Step 3: Convert and save model to NeMo format\n",
    "    output_file = os.path.join(cwd, 'Llama-3.2-1B.nemo')\n",
    "    logging.info(f\" Importing from {hf_source} → {output_file}...\")\n",
    "\n",
    "    llm.import_ckpt(\n",
    "        model=model_config,\n",
    "        source=hf_source,\n",
    "        output_path=output_file,\n",
    "    )\n",
    "\n",
    "    # Step 4: Confirm success\n",
    "    logging.info(f\" Done. Checkpoint saved to {os.path.abspath(output_file)}\")\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9d6ee7c6-23cd-43ff-9953-1a9f70e13df6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!huggingface-cli login --token \"hf_*******************************\"\n",
    "#!torchrun import_e5_large.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d093b29-a978-444d-ad5f-947e4e0c4b92",
   "metadata": {},
   "source": [
    "# Setup environment and training inputs  \n",
    "\n",
    "This cell:  \n",
    "- Imports required libraries (`os`, `nemo_run`, and NeMo’s `llm`).  \n",
    "- Defines key paths for training and experimentation:  \n",
    "  - `TRAIN_DATA_PATH`: Path to the dataset in **CustomRetrievalDataModule-compatible JSON** (prepared previously from AllNLI triplets).  \n",
    "  - `PRETRAINED_NEMO_MODEL`: Path to the converted E5-Large NeMo checkpoint (`.nemo` file imported earlier).  \n",
    "- Sets `NEMO_HOME` to the current working directory — NeMo will use this as its default workspace for logs, caches, and outputs.  \n",
    "\n",
    "Prerequisites:  \n",
    "- You should already have run:  \n",
    "  - The dataset export (`allnli_triplet.json`).  \n",
    "  - The Hugging Face → NeMo import for E5-Large (`e5-large-v2.nemo`).  \n",
    "- Ensure `nemo_run` (MCP runtime/launcher) and `nemo-toolkit` with `collections.llm` are installed.  \n",
    "\n",
    "Next steps:  \n",
    "- Attach this environment setup with your **training or finetuning script** for embedding retrieval tasks.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ec592999-b798-4f98-b23d-07173a532a21",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2025-08-26 13:42:48 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "      from .autonotebook import tqdm as notebook_tqdm\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:03 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/apex/transformer/tensor_parallel/cross_entropy.py:78: SyntaxWarning: invalid escape sequence '\\s'\n",
      "      \"\"\"\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:03 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py:49: SyntaxWarning: invalid escape sequence '\\_'\n",
      "      \"\"\"Run interleaved 1F1B schedule with communication between pipeline stages as needed.\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:03 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py:261: SyntaxWarning: invalid escape sequence '\\_'\n",
      "      \"\"\"Run non-interleaved 1F1B schedule, with communication between pipeline stages.\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:04 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/apex/contrib/optimizers/distributed_fused_adam.py:274: SyntaxWarning: invalid escape sequence '\\:'\n",
      "      \"\"\"Adam optimizer with ZeRO algorithm.\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:06 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pulp/tests/test_pulp.py:1776: SyntaxWarning: invalid escape sequence '\\d'\n",
      "      \"\"\"\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:06 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pulp/tests/test_pulp.py:1950: SyntaxWarning: invalid escape sequence '\\d'\n",
      "      command_line, option=\"strong\", grp_pattern=\"\\d+\"\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:19 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/jieba/__init__.py:44: SyntaxWarning: invalid escape sequence '\\.'\n",
      "      re_han_default = re.compile(\"([\\u4E00-\\u9FD5a-zA-Z0-9+#&\\._%\\-]+)\", re.U)\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:19 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/jieba/__init__.py:46: SyntaxWarning: invalid escape sequence '\\s'\n",
      "      re_skip_default = re.compile(\"(\\r\\n|\\s)\", re.U)\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:19 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/jieba/finalseg/__init__.py:78: SyntaxWarning: invalid escape sequence '\\.'\n",
      "      re_skip = re.compile(\"([a-zA-Z0-9]+(?:\\.\\d+)?%?)\")\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:22 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/lhotse/recipes/iwslt22_ta.py:323: SyntaxWarning: invalid escape sequence '\\s'\n",
      "      text = re.sub(\"\\s+\", \" \", text)\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:22 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/lhotse/recipes/iwslt22_ta.py:324: SyntaxWarning: invalid escape sequence '\\s'\n",
      "      text = re.sub(\"\\s+\\.\\s+\", \".\", text)\n",
      "    \n",
      "WARNING:root:We couldn't create lhotse utilities directory: /root/.lhotse/tools (not enough space/no permissions?)\n",
      "[NeMo W 2025-08-26 13:43:23 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:300: SyntaxWarning: invalid escape sequence '\\('\n",
      "      m = re.match('([su]([0-9]{1,2})p?) \\(([0-9]{1,2}) bit\\)$', token)\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:23 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:301: SyntaxWarning: invalid escape sequence '\\('\n",
      "      m2 = re.match('([su]([0-9]{1,2})p?)( \\(default\\))?$', token)\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:23 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:310: SyntaxWarning: invalid escape sequence '\\('\n",
      "      elif re.match('(flt)p?( \\(default\\))?$', token):\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:23 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:314: SyntaxWarning: invalid escape sequence '\\('\n",
      "      elif re.match('(dbl)p?( \\(default\\))?$', token):\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:23 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:170: RuntimeWarning: Couldn't find ffmpeg or avconv - defaulting to ffmpeg, but may not work\n",
      "      warn(\"Couldn't find ffmpeg or avconv - defaulting to ffmpeg, but may not work\", RuntimeWarning)\n",
      "    \n",
      "[NeMo W 2025-08-26 13:43:23 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/nvidia/dali/types.py:433: SyntaxWarning: invalid escape sequence '\\.'\n",
      "      _cupy_array_type_regex = re.compile(\".*cupy.*\\..*ndarray.*\")  # noqa: W605\n",
      "    \n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import nemo_run as run\n",
    "from nemo.collections import llm\n",
    "\n",
    "# Dataset path (downloaded separately)\n",
    "TRAIN_DATA_PATH = \"allnli_triplet.json\"\n",
    "\n",
    "# Pretrained E5 checkpoint (converted earlier using import_e5_large.py)\n",
    "PRETRAINED_NEMO_MODEL = \"e5-large-v2.nemo\"\n",
    "\n",
    "# NeMo working directory\n",
    "os.environ[\"NEMO_HOME\"] = os.getcwd()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fcdf26b-2796-4370-8954-e31c72da58e6",
   "metadata": {},
   "source": [
    "# Define Custom DataLoader for Triplet Retrieval  \n",
    "\n",
    "This helper function builds a **CustomRetrievalDataModule** that feeds triplet-style data (query, positive doc, negative doc) into the E5 embedding model during training.  \n",
    "\n",
    "Key points:  \n",
    "- **Inputs:**  \n",
    "  - `data_path`: Path to JSON dataset (e.g., `allnli_triplet.json`).  \n",
    "  - `dataset_identifier`: A label for experiment tracking (`'allnli_e5_triplet'`).  \n",
    "  - `seq_length`: Maximum token length for inputs (default: 512).  \n",
    "  - `micro_batch_size`: Per-GPU batch size (default: 16).  \n",
    "  - `global_batch_size`: Total batch size across GPUs (default: 64).  \n",
    "  - `tokenizer`: Tokenizer object (should match the pretrained `e5-large-v2` checkpoint).  \n",
    "  - `num_workers`: Dataloader workers for efficient loading (default: 8).  \n",
    "\n",
    "- **Triplet Mapping:**  \n",
    "  - `query_key=\"query\"` → input query sentence.  \n",
    "  - `pos_doc_key=\"pos_doc\"` → positive sentence (entailment).  \n",
    "  - `neg_doc_key=\"neg_doc\"` → negative sentence (contradiction/neutral).  \n",
    "\n",
    "- **Return:**  \n",
    "  - A configured `run.Config` object, wrapping `llm.CustomRetrievalDataModule`, which NeMo can use directly in training pipelines.  \n",
    "\n",
    "Usage example:  \n",
    "```python\n",
    "tokenizer = llm.AutoTokenizer.from_pretrained(\"hf://intfloat/e5-large-v2\")\n",
    "train_dataloader = get_custom_dataloader(\n",
    "    data_path=TRAIN_DATA_PATH,\n",
    "    tokenizer=tokenizer\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "db44e127-042e-46b0-b2a0-bacdd2790ebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_custom_dataloader(\n",
    "    data_path,\n",
    "    dataset_identifier='allnli_e5_triplet',\n",
    "    seq_length=512,\n",
    "    micro_batch_size=16,\n",
    "    global_batch_size=64,\n",
    "    tokenizer=None,\n",
    "    num_workers=8\n",
    "):\n",
    "    \"\"\"\n",
    "    Creates a CustomRetrievalDataModule for triplet training with E5.\n",
    "    \"\"\"\n",
    "    return run.Config(\n",
    "        llm.CustomRetrievalDataModule,\n",
    "        data_root=data_path,\n",
    "        dataset_identifier=dataset_identifier,\n",
    "        seq_length=seq_length,\n",
    "        micro_batch_size=micro_batch_size,\n",
    "        global_batch_size=global_batch_size,\n",
    "        tokenizer=tokenizer,\n",
    "        num_workers=num_workers,\n",
    "        query_key=\"query\",\n",
    "        pos_doc_key=\"pos_doc\",\n",
    "        neg_doc_key=\"neg_doc\",\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "974bfba5-29ce-46a6-bdb7-2f16455da96b",
   "metadata": {},
   "source": [
    "# Fine-tune E5-Large on AllNLI Triplets  \n",
    "\n",
    "This function defines and launches a **fine-tuning run** for the `e5-large-v2` embedding model on the **AllNLI triplet dataset**. It integrates the dataset, pretrained checkpoint, and NeMo recipe into a runnable training workflow.  \n",
    "\n",
    "### Workflow:\n",
    "1. **Pretrained checkpoint**  \n",
    "   - Loads `e5-large-v2.nemo` converted earlier with `import_e5_large.py`.  \n",
    "\n",
    "2. **Dataset integration**  \n",
    "   - Calls `get_custom_dataloader(...)` to wrap the `allnli_triplet.json` dataset in a `CustomRetrievalDataModule`.  \n",
    "\n",
    "3. **Training recipe**  \n",
    "   - Loads NeMo’s built-in `e5_340m.finetune_recipe` as a base.  \n",
    "   - Overrides settings for this run (e.g. model path, dataset, hardware params).  \n",
    "   - Uses **single-GPU, single-node setup** (modifiable).  \n",
    "\n",
    "4. **Customized hyperparameters**  \n",
    "   - Enables **global in-batch negatives** for contrastive retrieval.  \n",
    "   - Learning rate: `5e-6` (with scheduler min LR = `5e-7`).  \n",
    "   - Training duration: 100 steps (toy example, extend as needed).  \n",
    "   - Validation every 10 steps, checking 5 mini-batches each time.  \n",
    "\n",
    "5. **Execution**  \n",
    "   - `run.run(recipe, executor=run.LocalExecutor())` launches training locally.  \n",
    "\n",
    "### Usage:\n",
    "```python\n",
    "train_e5_on_allnli(TRAIN_DATA_PATH)\n",
    "```\n",
    "\n",
    "This will:  \n",
    "- Load your pretrained **E5-Large `.nemo` checkpoint**  \n",
    "- Fine-tune it on **AllNLI triplets**  \n",
    "- Save outputs/logs inside your current `NEMO_HOME` workspace  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e39db94d-d10c-4756-8367-0935f03be015",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_e5_on_allnli(json_file_path):\n",
    "    \"\"\"\n",
    "    Fine-tune the E5 model on the AllNLI triplet dataset.\n",
    "    \"\"\"\n",
    "    pretrained_model_path = os.path.abspath(PRETRAINED_NEMO_MODEL)\n",
    "\n",
    "    # Create datamodule\n",
    "    datamodule = get_custom_dataloader(\n",
    "        data_path=json_file_path,\n",
    "        dataset_identifier=\"allnli_e5_triplet\"\n",
    "    )\n",
    "\n",
    "    # Load recipe\n",
    "    recipe = llm.recipes.e5_340m.finetune_recipe(\n",
    "        name=\"allnli_e5_large_finetune\",\n",
    "        resume_path=pretrained_model_path,\n",
    "        num_nodes=1,\n",
    "        num_gpus_per_node=1,\n",
    "    )\n",
    "\n",
    "    # Customize recipe params\n",
    "    recipe.model.config.global_in_batch_negatives = True\n",
    "    recipe.optim.config.lr = 5e-6\n",
    "    recipe.optim.lr_scheduler.min_lr = 5e-7\n",
    "    recipe.trainer.max_steps = 100\n",
    "    recipe.trainer.val_check_interval = 10\n",
    "    recipe.trainer.limit_val_batches = 5\n",
    "    recipe.data = datamodule\n",
    "\n",
    "    # Run training\n",
    "    run.run(recipe, executor=run.LocalExecutor())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d3ab4b4-52c4-46e2-8b57-7f50d9bf1a4b",
   "metadata": {},
   "source": [
    "# Launch Fine-tuning Run  \n",
    "\n",
    "This cell executes the full **fine-tuning job** by calling:  \n",
    "\n",
    "```python\n",
    "train_e5_on_allnli(TRAIN_DATA_PATH)\n",
    "```\n",
    "\n",
    "What happens when run:  \n",
    "- Loads the **pretrained `e5-large-v2.nemo` checkpoint**.  \n",
    "- Prepares the **AllNLI triplet dataset** (`allnli_triplet.json`) via the custom dataloader.  \n",
    "- Builds and configures the **fine-tune recipe** (single GPU, local execution).  \n",
    "- Starts training with the specified hyperparameters (100 steps, periodic validation).  \n",
    "- Saves logs, checkpoints, and artifacts inside your `NEMO_HOME` directory.  \n",
    "\n",
    "⚠️ Note: This configuration is set up for a **short debugging/trial run** (100 steps).  \n",
    "- Increase `recipe.trainer.max_steps` and adjust validation intervals for full training.  \n",
    "- For multi-GPU or cluster training, replace `LocalExecutor()` with a distributed executor (or use `torchrun`).  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b2266012-0b14-48e0-a7f9-b63371261da8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #00ff00; text-decoration-color: #00ff00\">─── </span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">Entering Experiment nemo.collections.llm.api.finetune with id: nemo.collections.llm.api.finetune_1756196008</span><span style=\"color: #00ff00; text-decoration-color: #00ff00\"> ───</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[92m─── \u001b[0m\u001b[1;35mEntering Experiment nemo.collections.llm.api.finetune with id: nemo.collections.llm.api.finetune_1756196008\u001b[0m\u001b[92m ───\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Log directory is: /root/.nemo_run/experiments/nemo.collections.llm.api.finetune/nemo.collections.llm.api.finetune_1756196008/nemo.collections.llm.api.finetune\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[13:43:28] </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">Launching job nemo.collections.llm.api.finetune for experiment </span>                        <a href=\"file:///opt/Run/nemo_run/run/experiment.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">experiment.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///opt/Run/nemo_run/run/experiment.py#744\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">744</span></a>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">           </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">nemo.collections.llm.api.finetune</span>                                                      <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                 </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[13:43:28]\u001b[0m\u001b[2;36m \u001b[0m\u001b[1;36mLaunching job nemo.collections.llm.api.finetune for experiment \u001b[0m                        \u001b]8;id=550310;file:///opt/Run/nemo_run/run/experiment.py\u001b\\\u001b[2mexperiment.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=771114;file:///opt/Run/nemo_run/run/experiment.py#744\u001b\\\u001b[2m744\u001b[0m\u001b]8;;\u001b\\\n",
       "\u001b[2;36m           \u001b[0m\u001b[1;36mnemo.collections.llm.api.finetune\u001b[0m                                                      \u001b[2m                 \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Log directory is: /root/.nemo_run/experiments/nemo.collections.llm.api.finetune/nemo.collections.llm.api.finetune_1756196008/nemo.collections.llm.api.finetune\n",
      "Launched app: local_persistent://nemo_run/nemo.collections.llm.api.finetune-zn4250k9j945t\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #00ff00; text-decoration-color: #00ff00\">────────────────── </span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">Waiting for Experiment nemo.collections.llm.api.finetune_1756196008 to finish</span><span style=\"color: #00ff00; text-decoration-color: #00ff00\"> ──────────────────</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[92m────────────────── \u001b[0m\u001b[1;35mWaiting for Experiment nemo.collections.llm.api.finetune_1756196008 to finish\u001b[0m\u001b[92m ──────────────────\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">Experiment Status for</span> <span style=\"color: #ffaf00; text-decoration-color: #ffaf00; font-weight: bold\">nemo.collections.llm.api.finetune_1756196008</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;32mExperiment Status for\u001b[0m \u001b[1;38;5;214mnemo.collections.llm.api.finetune_1756196008\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "<span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">Task 0</span>: <span style=\"color: #ffaf00; text-decoration-color: #ffaf00; font-weight: bold\">nemo.collections.llm.api.finetune</span>\n",
       "- <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">Status</span>: RUNNING\n",
       "- <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">Executor</span>: LocalExecutor\n",
       "- <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">Job id</span>: nemo.collections.llm.api.finetune-zn4250k9j945t\n",
       "- <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">Local Directory</span>: /root/.nemo_run/experiments/nemo.collections.llm.api.finetune/nemo.collections.llm.api.finetune_1756196008/nemo.collections.llm.api.finetune\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n",
       "\u001b[1;32mTask 0\u001b[0m: \u001b[1;38;5;214mnemo.collections.llm.api.finetune\u001b[0m\n",
       "- \u001b[1;32mStatus\u001b[0m: RUNNING\n",
       "- \u001b[1;32mExecutor\u001b[0m: LocalExecutor\n",
       "- \u001b[1;32mJob id\u001b[0m: nemo.collections.llm.api.finetune-zn4250k9j945t\n",
       "- \u001b[1;32mLocal Directory\u001b[0m: /root/.nemo_run/experiments/nemo.collections.llm.api.finetune/nemo.collections.llm.api.finetune_1756196008/nemo.collections.llm.api.finetune\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Waiting for job nemo.collections.llm.api.finetune-zn4250k9j945t to finish [log=True]...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i.finetune/0 [NeMo W 2025-08-26 13:43:58 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/apex/transformer/tensor_parallel/cross_entropy.py:78: SyntaxWarning: invalid escape sequence '\\s'\n",
      "i.finetune/0       \"\"\"\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:43:58 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py:49: SyntaxWarning: invalid escape sequence '\\_'\n",
      "i.finetune/0       \"\"\"Run interleaved 1F1B schedule with communication between pipeline stages as needed.\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:43:58 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py:261: SyntaxWarning: invalid escape sequence '\\_'\n",
      "i.finetune/0       \"\"\"Run non-interleaved 1F1B schedule, with communication between pipeline stages.\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:43:58 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/apex/contrib/optimizers/distributed_fused_adam.py:274: SyntaxWarning: invalid escape sequence '\\:'\n",
      "i.finetune/0       \"\"\"Adam optimizer with ZeRO algorithm.\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:43:59 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pulp/tests/test_pulp.py:1776: SyntaxWarning: invalid escape sequence '\\d'\n",
      "i.finetune/0       \"\"\"\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:43:59 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pulp/tests/test_pulp.py:1950: SyntaxWarning: invalid escape sequence '\\d'\n",
      "i.finetune/0       command_line, option=\"strong\", grp_pattern=\"\\d+\"\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:11 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/jieba/__init__.py:44: SyntaxWarning: invalid escape sequence '\\.'\n",
      "i.finetune/0       re_han_default = re.compile(\"([\\u4E00-\\u9FD5a-zA-Z0-9+#&\\._%\\-]+)\", re.U)\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:11 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/jieba/__init__.py:46: SyntaxWarning: invalid escape sequence '\\s'\n",
      "i.finetune/0       re_skip_default = re.compile(\"(\\r\\n|\\s)\", re.U)\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:11 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/jieba/finalseg/__init__.py:78: SyntaxWarning: invalid escape sequence '\\.'\n",
      "i.finetune/0       re_skip = re.compile(\"([a-zA-Z0-9]+(?:\\.\\d+)?%?)\")\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:13 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/lhotse/recipes/iwslt22_ta.py:323: SyntaxWarning: invalid escape sequence '\\s'\n",
      "i.finetune/0       text = re.sub(\"\\s+\", \" \", text)\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:13 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/lhotse/recipes/iwslt22_ta.py:324: SyntaxWarning: invalid escape sequence '\\s'\n",
      "i.finetune/0       text = re.sub(\"\\s+\\.\\s+\", \".\", text)\n",
      "i.finetune/0     \n",
      "i.finetune/0 WARNING:root:We couldn't create lhotse utilities directory: /root/.lhotse/tools (not enough space/no permissions?)\n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:14 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:300: SyntaxWarning: invalid escape sequence '\\('\n",
      "i.finetune/0       m = re.match('([su]([0-9]{1,2})p?) \\(([0-9]{1,2}) bit\\)$', token)\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:14 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:301: SyntaxWarning: invalid escape sequence '\\('\n",
      "i.finetune/0       m2 = re.match('([su]([0-9]{1,2})p?)( \\(default\\))?$', token)\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:14 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:310: SyntaxWarning: invalid escape sequence '\\('\n",
      "i.finetune/0       elif re.match('(flt)p?( \\(default\\))?$', token):\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:14 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:314: SyntaxWarning: invalid escape sequence '\\('\n",
      "i.finetune/0       elif re.match('(dbl)p?( \\(default\\))?$', token):\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:14 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/pydub/utils.py:170: RuntimeWarning: Couldn't find ffmpeg or avconv - defaulting to ffmpeg, but may not work\n",
      "i.finetune/0       warn(\"Couldn't find ffmpeg or avconv - defaulting to ffmpeg, but may not work\", RuntimeWarning)\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:16 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/nvidia/dali/types.py:433: SyntaxWarning: invalid escape sequence '\\.'\n",
      "i.finetune/0       _cupy_array_type_regex = re.compile(\".*cupy.*\\..*ndarray.*\")  # noqa: W605\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:19 nemo_logging:393] data_root: ['allnli_triplet.json'] will be split to 0.95:0.04:0.01 used for train/val/test\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:20 nemo_logging:393] Experiments will be logged at /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20\n",
      "i.finetune/0 INFO:pytorch_lightning.utilities.rank_zero:GPU available: True (cuda), used: True\n",
      "i.finetune/0 INFO:pytorch_lightning.utilities.rank_zero:TPU available: False, using: 0 TPU cores\n",
      "i.finetune/0 INFO:pytorch_lightning.utilities.rank_zero:HPU available: False, using: 0 HPUs\n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:20 nemo_logging:405] \"update_logger_directory\" is True. Overwriting tensorboard logger \"save_dir\" to /host_pwd/PythonNotebook/nemo_experiments/tb_logs\n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:20 nemo_logging:405] The Trainer already contains a ModelCheckpoint callback. This will be overwritten.\n",
      "i.finetune/0 [NeMo W 2025-08-26 13:44:20 nemo_logging:405] The checkpoint callback was told to monitor a validation value and trainer's max_steps was set to 100. Please ensure that max_steps will run for at least 1 epochs to ensure that checkpointing will not error out.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Fixing mis-match between ddp-config & mcore-optimizer config\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Rank 0 has data parallel group : [0]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Rank 0 has combined group of data parallel and context parallel : [0]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] All data parallel group ranks with context parallel combined: [[0]]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Ranks 0 has data parallel rank: 0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Rank 0 has context parallel group: [0]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] All context parallel group ranks: [[0]]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Ranks 0 has context parallel rank: 0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Rank 0 has model parallel group: [0]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] All model parallel group ranks: [[0]]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Rank 0 has tensor model parallel group: [0]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] All tensor model parallel group ranks: [[0]]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Rank 0 has tensor model parallel rank: 0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Rank 0 has pipeline model parallel group: [0]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Rank 0 has embedding group: [0]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] All pipeline model parallel group ranks: [[0]]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Rank 0 has pipeline model parallel rank 0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] All embedding group ranks: [[0]]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Rank 0 has embedding rank: 0\n",
      "i.finetune/0 INFO:pytorch_lightning.utilities.rank_zero:----------------------------------------------------------------------------------------------------\n",
      "i.finetune/0 distributed_backend=nccl\n",
      "i.finetune/0 All distributed processes registered. Starting with 1 processes\n",
      "i.finetune/0 ----------------------------------------------------------------------------------------------------\n",
      "i.finetune/0 \n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:21 nemo_logging:393] Preprocessing CustomRetrievalDataModule to jsonl format and splitting...\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:56 nemo_logging:393] training split saved to /host_pwd/PythonNotebook/datasets/allnli_e5_triplet/training.jsonl\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:57 nemo_logging:393] validation split saved to /host_pwd/PythonNotebook/datasets/allnli_e5_triplet/validation.jsonl\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:57 nemo_logging:393] test split saved to /host_pwd/PythonNotebook/datasets/allnli_e5_triplet/test.jsonl\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:57 nemo_logging:393] Padded vocab_size: 30592, original vocab_size: 30522, dummy tokens: 70.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:58 nemo_logging:393] Copying Trainer's 'max_steps' (100) to LR scheduler's 'max_steps'.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:58 num_microbatches_calculator:228] setting number of microbatches to constant 4\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:58 nemo_logging:393]  > number of parameters on (tensor, pipeline) model parallel rank (0 ,0): 335244160\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:58 utils:532] Setting up DistributedDataParallel with config DistributedDataParallelConfig(grad_reduce_in_fp32=True, overlap_grad_reduce=False, overlap_param_gather=False, align_param_gather=False, use_distributed_optimizer=True, num_distributed_optimizer_instances=1, check_for_nan_in_grad=True, check_for_large_grads=False, bucket_size=None, pad_buckets_for_high_nccl_busbw=False, average_in_collective=False, fp8_param_gather=False, use_custom_fsdp=False, data_parallel_sharding_strategy='no_shard', gradient_reduce_div_fusion=True, suggested_communication_unit_size=None, preserve_fp32_weights=True, keep_fp8_transpose_cache_when_using_custom_fsdp=False)\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:58 utils:553] Number of buckets for gradient all-reduce / reduce-scatter: 1\n",
      "i.finetune/0     Params for bucket 1 (335244160 elements, 335244160 padded size):\n",
      "i.finetune/0     \tmodule.encoder.layers.23.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.20.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.16.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.14.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.11.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.8.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.6.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.3.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.0.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.23.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.21.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.18.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.14.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.12.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.9.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.6.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.4.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.1.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.23.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.19.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.15.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.14.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.10.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.7.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.6.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.2.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.embedding.position_embeddings.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.23.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.20.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.18.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.16.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.14.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.11.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.8.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.6.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.3.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.1.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.22.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.19.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.15.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.13.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.10.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.7.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.5.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.2.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.embedding.tokentype_embeddings.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.0.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.22.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.20.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.17.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.16.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.13.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.11.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.8.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.5.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.3.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.0.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.23.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.22.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.18.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.14.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.13.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.9.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.6.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.5.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.1.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.22.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.19.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.15.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.13.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.10.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.7.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.5.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.2.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.23.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.21.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.18.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.17.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.14.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.12.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.9.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.6.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.4.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.1.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.21.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.19.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.15.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.12.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.10.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.7.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.4.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.2.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.22.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.21.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.17.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.16.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.13.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.12.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.8.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.5.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.4.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.0.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.18.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.23.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.21.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.14.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.12.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.9.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.6.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.4.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.1.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.0.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.22.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.20.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.16.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.13.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.11.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.8.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.5.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.3.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.23.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.20.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.18.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.17.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.14.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.11.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.9.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.6.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.3.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.1.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.21.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.20.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.15.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.12.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.11.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.7.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.4.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.3.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.0.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.22.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.20.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.17.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.16.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.13.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.11.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.8.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.5.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.3.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.0.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.21.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.19.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.15.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.12.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.10.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.7.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.4.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.2.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.embedding.word_embeddings.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.22.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.19.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.16.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.13.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.10.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.8.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.5.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.2.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.23.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.20.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.19.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.18.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.14.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.11.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.10.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.6.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.3.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.2.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.0.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.21.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.19.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.15.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.12.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.10.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.7.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.4.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.2.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.23.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.20.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.17.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.17.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.14.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.11.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.9.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.6.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.3.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.1.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.initial_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.21.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.18.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.15.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.12.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.9.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.7.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.4.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.1.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.22.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.19.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.17.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.13.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.10.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.9.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.5.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.2.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.23.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.20.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.18.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.17.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.14.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.11.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.9.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.6.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.3.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.1.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.22.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.19.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.16.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.13.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.10.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.8.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.5.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.2.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.0.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.output_layer.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.23.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.20.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.17.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.16.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.14.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.11.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.8.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.6.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.3.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.0.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.pooler.dense.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.21.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.18.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.16.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.12.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.9.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.8.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.4.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.1.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.pooler.dense.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.22.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.19.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.16.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.13.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.10.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.8.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.5.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.2.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.initial_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.21.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.18.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.17.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.15.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.12.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.9.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.7.post_mlp_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.4.self_attention.linear_proj.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.1.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.22.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.19.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.15.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.13.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.10.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.7.post_mlp_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.5.mlp.linear_fc1.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.2.self_attention.linear_proj.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.0.mlp.linear_fc2.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.20.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.17.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.16.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.15.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.11.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.8.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.7.post_att_layernorm.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.3.self_attention.linear_qkv.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.0.mlp.linear_fc2.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.21.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.18.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.15.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.12.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.9.post_att_layernorm.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.7.self_attention.linear_qkv.bias\n",
      "i.finetune/0     \tmodule.encoder.layers.4.mlp.linear_fc1.weight\n",
      "i.finetune/0     \tmodule.encoder.layers.1.post_att_layernorm.weight\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:58 utils:532] Setting up optimizer with config OptimizerConfig(optimizer='adam', lr=5e-06, min_lr=None, decoupled_lr=None, decoupled_min_lr=None, weight_decay=0.1, fp16=False, bf16=True, params_dtype=torch.bfloat16, use_precision_aware_optimizer=False, main_grads_dtype=torch.float32, main_params_dtype=torch.float32, exp_avg_dtype=torch.float32, exp_avg_sq_dtype=torch.float32, loss_scale=None, initial_loss_scale=4294967296, min_loss_scale=1.0, loss_scale_window=1000, hysteresis=2, adam_beta1=0.9, adam_beta2=0.98, adam_eps=1e-05, sgd_momentum=0.9, use_distributed_optimizer=True, overlap_param_gather_with_optimizer_step=False, optimizer_cpu_offload=False, optimizer_offload_fraction=0.0, use_torch_optimizer_for_cpu_offload=False, overlap_cpu_optimizer_d2h_h2d=False, pin_cpu_grads=True, pin_cpu_params=True, clip_grad=1.0, log_num_zeros_in_grad=False, barrier_with_L1_time=False, timers=None, config_logger_dir='')\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:58 nemo_logging:393] Doing selective restore from RestoreConfig(path='/host_pwd/PythonNotebook/e5-large-v2.nemo', adapter_path=None, load_model_state=True, load_optim_state=False, load_artifacts=True)\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:44:58 nemo_logging:393] Using <megatron.core.dist_checkpointing.strategies.fully_parallel.FullyParallelLoadStrategyWrapper object at 0x1552788e3b60> dist-ckpt load strategy.\n",
      "i.finetune/0 LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:01 nemo_logging:393] Global Checkpoint Load : Rank : 0 : Start time : 1756196098.856s : Time spent in load_checkpoint: 2.614s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:01 nemo_logging:393] Restoring model weights from RestoreConfig(path='/host_pwd/PythonNotebook/e5-large-v2.nemo', adapter_path=None, load_model_state=True, load_optim_state=False, load_artifacts=True)\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:01 nemo_logging:393] Finished restoring from RestoreConfig(path='/host_pwd/PythonNotebook/e5-large-v2.nemo', adapter_path=None, load_model_state=True, load_optim_state=False, load_artifacts=True), cleaning up.\n",
      "i.finetune/0 \n",
      "i.finetune/0   | Name           | Type              | Params | Mode \n",
      "i.finetune/0 -------------------------------------------------------------\n",
      "i.finetune/0 0 | module         | DDP               | 335 M  | train\n",
      "i.finetune/0 1 | embedding_head | BertEmbeddingHead | 0      | train\n",
      "i.finetune/0 -------------------------------------------------------------\n",
      "i.finetune/0 335 M     Trainable params\n",
      "i.finetune/0 0         Non-trainable params\n",
      "i.finetune/0 335 M     Total params\n",
      "i.finetune/0 1,340.977 Total estimated model params size (MB)\n",
      "i.finetune/0 543       Modules in train mode\n",
      "i.finetune/0 0         Modules in eval mode\n",
      "i.finetune/0  > WARNING: could not find index map file /host_pwd/PythonNotebook/datasets/allnli_e5_triplet/validation.jsonl_validation.jsonl_indexmap_6432mns_510msl_0.00ssp_1234s.npy, building the indices on rank 0 ...\n",
      "i.finetune/0     using uint32 for data mapping...\n",
      "i.finetune/0     using:\n",
      "i.finetune/0      number of documents:            22314\n",
      "i.finetune/0      sentences range:                [0, 22314)\n",
      "i.finetune/0      total number of sentences:      22314\n",
      "i.finetune/0      number of epochs:               2147483646\n",
      "i.finetune/0      maximum number of samples:      6432\n",
      "i.finetune/0      maximum sequence length:        510\n",
      "i.finetune/0      short sequence probability:     0\n",
      "i.finetune/0      short sequence ration (1/prob): 0\n",
      "i.finetune/0      seed:                           1234\n",
      "i.finetune/0     reached 6432 samples after 1 epochs ...\n",
      "i.finetune/0    number of empty documents: 0\n",
      "i.finetune/0    number of documents with one sentence: 22314\n",
      "i.finetune/0    number of documents with long sentences: 0\n",
      "i.finetune/0    will create mapping for 22314 samples\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:02 nemo_logging:393] Creating EmbeddingDataset with seed=1234,\n",
      "i.finetune/0     add_bos=False, add_eos=True,\n",
      "i.finetune/0     max_seq_length=512, min_seq_length=1,\n",
      "i.finetune/0     pad_token_id=102, negative_sample_strategy=first,\n",
      "i.finetune/0     num_hard_negatives=1.\n",
      "i.finetune/0 [NeMo W 2025-08-26 13:45:15 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:431: It is recommended to use `self.log('global_batch_size', ..., sync_dist=True)` when logging on epoch level in distributed setting to accumulate the metric across devices.\n",
      "i.finetune/0     \n",
      "i.finetune/0 [NeMo W 2025-08-26 13:45:15 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:431: It is recommended to use `self.log('val_loss', ..., sync_dist=True)` when logging on epoch level in distributed setting to accumulate the metric across devices.\n",
      "i.finetune/0     \n",
      "i.finetune/0  > WARNING: could not find index map file /host_pwd/PythonNotebook/datasets/allnli_e5_triplet/training.jsonl_training.jsonl_indexmap_6432mns_510msl_0.00ssp_1234s.npy, building the indices on rank 0 ...\n",
      "i.finetune/0     using uint32 for data mapping...\n",
      "i.finetune/0     using:\n",
      "i.finetune/0      number of documents:            529957\n",
      "i.finetune/0      sentences range:                [0, 529957)\n",
      "i.finetune/0      total number of sentences:      529957\n",
      "i.finetune/0      number of epochs:               2147483646\n",
      "i.finetune/0      maximum number of samples:      6432\n",
      "i.finetune/0      maximum sequence length:        510\n",
      "i.finetune/0      short sequence probability:     0\n",
      "i.finetune/0      short sequence ration (1/prob): 0\n",
      "i.finetune/0      seed:                           1234\n",
      "i.finetune/0     reached 6432 samples after 1 epochs ...\n",
      "i.finetune/0    number of empty documents: 0\n",
      "i.finetune/0    number of documents with one sentence: 529957\n",
      "i.finetune/0    number of documents with long sentences: 0\n",
      "i.finetune/0    will create mapping for 529957 samples\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:15 nemo_logging:393] Creating EmbeddingDataset with seed=1234,\n",
      "i.finetune/0     add_bos=False, add_eos=True,\n",
      "i.finetune/0     max_seq_length=512, min_seq_length=1,\n",
      "i.finetune/0     pad_token_id=102, negative_sample_strategy=first,\n",
      "i.finetune/0     num_hard_negatives=1.\n",
      "i.finetune/0 [NeMo W 2025-08-26 13:45:20 rerun_state_machine:1264] Implicit initialization of Rerun State Machine!\n",
      "i.finetune/0 [NeMo W 2025-08-26 13:45:20 rerun_state_machine:239] RerunStateMachine initialized in mode RerunMode.DISABLED\n",
      "i.finetune/0 Training epoch 0, iteration 0/99 | lr: 9.804e-08 | global_batch_size: 64 | global_step: 0 | reduced_train_loss: 1.077 | train_step_timing in s: 4.49\n",
      "i.finetune/0 Training epoch 0, iteration 1/99 | lr: 1.961e-07 | global_batch_size: 64 | global_step: 1 | reduced_train_loss: 1.078 | train_step_timing in s: 2.744 | consumed_samples: 128\n",
      "i.finetune/0 Training epoch 0, iteration 2/99 | lr: 2.941e-07 | global_batch_size: 64 | global_step: 2 | reduced_train_loss: 0.999 | train_step_timing in s: 0.5949 | consumed_samples: 192\n",
      "i.finetune/0 Training epoch 0, iteration 3/99 | lr: 3.922e-07 | global_batch_size: 64 | global_step: 3 | reduced_train_loss: 0.9634 | train_step_timing in s: 0.7147 | consumed_samples: 256\n",
      "i.finetune/0 Training epoch 0, iteration 4/99 | lr: 4.902e-07 | global_batch_size: 64 | global_step: 4 | reduced_train_loss: 1.066 | train_step_timing in s: 0.5878 | consumed_samples: 320\n",
      "i.finetune/0 Training epoch 0, iteration 5/99 | lr: 5.882e-07 | global_batch_size: 64 | global_step: 5 | reduced_train_loss: 0.928 | train_step_timing in s: 1.206 | consumed_samples: 384\n",
      "i.finetune/0 Training epoch 0, iteration 6/99 | lr: 6.863e-07 | global_batch_size: 64 | global_step: 6 | reduced_train_loss: 0.9299 | train_step_timing in s: 0.5924 | consumed_samples: 448\n",
      "i.finetune/0 Training epoch 0, iteration 7/99 | lr: 7.843e-07 | global_batch_size: 64 | global_step: 7 | reduced_train_loss: 0.8691 | train_step_timing in s: 0.599 | consumed_samples: 512\n",
      "i.finetune/0 Training epoch 0, iteration 8/99 | lr: 8.824e-07 | global_batch_size: 64 | global_step: 8 | reduced_train_loss: 0.8999 | train_step_timing in s: 0.846 | consumed_samples: 576\n",
      "i.finetune/0 Training epoch 0, iteration 9/99 | lr: 9.804e-07 | global_batch_size: 64 | global_step: 9 | reduced_train_loss: 0.9607 | train_step_timing in s: 0.5919 | consumed_samples: 640\n",
      "i.finetune/0 Validation: iteration 1/0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:33 nemo_logging:393] Using FullyParallelSaveStrategyWrapper(torch_dist, 1) dist-ckpt save strategy.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:33 nemo_logging:393] Global Checkpoint Save : Rank: 0 : Iteration: 10 : Start time: 1756196133.682s : Save duration: 0.232s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:38 nemo_logging:393] Scheduled async checkpoint save for /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.74-step=9-consumed_samples=640.0-last.ckpt\n",
      "i.finetune/0 Training epoch 0, iteration 10/99 | lr: 1.078e-06 | global_batch_size: 64 | global_step: 10 | reduced_train_loss: 0.9127 | train_step_timing in s: 1.9 | consumed_samples: 704 | val_loss: 0.7394\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:40 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 11/99 | lr: 1.176e-06 | global_batch_size: 64 | global_step: 11 | reduced_train_loss: 1.089 | train_step_timing in s: 2.262 | consumed_samples: 768 | val_loss: 0.7394\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:43 nemo_logging:393] Async finalization time took 0.003 s\n",
      "i.finetune/0 Training epoch 0, iteration 12/99 | lr: 1.275e-06 | global_batch_size: 64 | global_step: 12 | reduced_train_loss: 0.9896 | train_step_timing in s: 2.389 | consumed_samples: 832 | val_loss: 0.7394\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:45 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 13/99 | lr: 1.373e-06 | global_batch_size: 64 | global_step: 13 | reduced_train_loss: 1.014 | train_step_timing in s: 1.287 | consumed_samples: 896 | val_loss: 0.7394\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:47 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 14/99 | lr: 1.471e-06 | global_batch_size: 64 | global_step: 14 | reduced_train_loss: 1.071 | train_step_timing in s: 1.298 | consumed_samples: 960 | val_loss: 0.7394\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:49 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 15/99 | lr: 1.569e-06 | global_batch_size: 64 | global_step: 15 | reduced_train_loss: 1.13 | train_step_timing in s: 0.6122 | consumed_samples: 1024 | val_loss: 0.7394\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:50 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 16/99 | lr: 1.667e-06 | global_batch_size: 64 | global_step: 16 | reduced_train_loss: 0.9134 | train_step_timing in s: 0.6112 | consumed_samples: 1088 | val_loss: 0.7394\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:51 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 17/99 | lr: 1.765e-06 | global_batch_size: 64 | global_step: 17 | reduced_train_loss: 0.8961 | train_step_timing in s: 0.8414 | consumed_samples: 1152 | val_loss: 0.7394\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:52 nemo_logging:393] Successfully saved checkpoint from iteration      10 to /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.74-step=9-consumed_samples=640.0-last.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:52 nemo_logging:393] Async checkpoint save for step 10 (/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.74-step=9-consumed_samples=640.0-last.ckpt) finalized successfully.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:52 nemo_logging:393] Async finalization time took 0.053 s\n",
      "i.finetune/0 Training epoch 0, iteration 18/99 | lr: 1.863e-06 | global_batch_size: 64 | global_step: 18 | reduced_train_loss: 0.8286 | train_step_timing in s: 0.6303 | consumed_samples: 1216 | val_loss: 0.7394\n",
      "i.finetune/0 Training epoch 0, iteration 19/99 | lr: 1.961e-06 | global_batch_size: 64 | global_step: 19 | reduced_train_loss: 0.9067 | train_step_timing in s: 0.8786 | consumed_samples: 1280 | val_loss: 0.7394\n",
      "i.finetune/0 Validation: iteration 1/0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:45:59 nemo_logging:393] Global Checkpoint Save : Rank: 0 : Iteration: 20 : Start time: 1756196159.043s : Save duration: 0.175s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:00 nemo_logging:393] Scheduled async checkpoint save for /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.66-step=19-consumed_samples=1280.0-last.ckpt\n",
      "i.finetune/0 Training epoch 0, iteration 20/99 | lr: 2.059e-06 | global_batch_size: 64 | global_step: 20 | reduced_train_loss: 0.7763 | train_step_timing in s: 1.857 | consumed_samples: 1344 | val_loss: 0.6562\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:02 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 21/99 | lr: 2.157e-06 | global_batch_size: 64 | global_step: 21 | reduced_train_loss: 0.8467 | train_step_timing in s: 2.682 | consumed_samples: 1408 | val_loss: 0.6562\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:05 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 22/99 | lr: 2.255e-06 | global_batch_size: 64 | global_step: 22 | reduced_train_loss: 0.7575 | train_step_timing in s: 1.819 | consumed_samples: 1472 | val_loss: 0.6562\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:07 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 23/99 | lr: 2.353e-06 | global_batch_size: 64 | global_step: 23 | reduced_train_loss: 0.8409 | train_step_timing in s: 1.247 | consumed_samples: 1536 | val_loss: 0.6562\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:09 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 24/99 | lr: 2.451e-06 | global_batch_size: 64 | global_step: 24 | reduced_train_loss: 0.8239 | train_step_timing in s: 0.7425 | consumed_samples: 1600 | val_loss: 0.6562\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:10 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 25/99 | lr: 2.549e-06 | global_batch_size: 64 | global_step: 25 | reduced_train_loss: 0.7847 | train_step_timing in s: 0.7011 | consumed_samples: 1664 | val_loss: 0.6562\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:12 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 26/99 | lr: 2.647e-06 | global_batch_size: 64 | global_step: 26 | reduced_train_loss: 0.7069 | train_step_timing in s: 0.761 | consumed_samples: 1728 | val_loss: 0.6562\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:13 nemo_logging:393] Successfully saved checkpoint from iteration      20 to /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.66-step=19-consumed_samples=1280.0-last.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:13 nemo_logging:393] Async checkpoint save for step 20 (/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.66-step=19-consumed_samples=1280.0-last.ckpt) finalized successfully.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:15 nemo_logging:393] Async finalization time took 2.112 s\n",
      "i.finetune/0 Training epoch 0, iteration 27/99 | lr: 2.745e-06 | global_batch_size: 64 | global_step: 27 | reduced_train_loss: 0.7356 | train_step_timing in s: 0.6045 | consumed_samples: 1792 | val_loss: 0.6562\n",
      "i.finetune/0 Training epoch 0, iteration 28/99 | lr: 2.843e-06 | global_batch_size: 64 | global_step: 28 | reduced_train_loss: 0.7111 | train_step_timing in s: 0.6273 | consumed_samples: 1856 | val_loss: 0.6562\n",
      "i.finetune/0 Training epoch 0, iteration 29/99 | lr: 2.941e-06 | global_batch_size: 64 | global_step: 29 | reduced_train_loss: 0.7725 | train_step_timing in s: 0.6023 | consumed_samples: 1920 | val_loss: 0.6562\n",
      "i.finetune/0 Validation: iteration 1/0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:21 nemo_logging:393] Global Checkpoint Save : Rank: 0 : Iteration: 30 : Start time: 1756196181.353s : Save duration: 0.180s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:21 nemo_logging:393] Scheduled async checkpoint save for /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.48-step=29-consumed_samples=1920.0-last.ckpt\n",
      "i.finetune/0 Training epoch 0, iteration 30/99 | lr: 3.039e-06 | global_batch_size: 64 | global_step: 30 | reduced_train_loss: 0.6742 | train_step_timing in s: 2.869 | consumed_samples: 1984 | val_loss: 0.477\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:24 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 31/99 | lr: 3.137e-06 | global_batch_size: 64 | global_step: 31 | reduced_train_loss: 0.6405 | train_step_timing in s: 2.175 | consumed_samples: 2048 | val_loss: 0.477\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:27 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 32/99 | lr: 3.235e-06 | global_batch_size: 64 | global_step: 32 | reduced_train_loss: 0.5062 | train_step_timing in s: 2.185 | consumed_samples: 2112 | val_loss: 0.477\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:29 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 33/99 | lr: 3.333e-06 | global_batch_size: 64 | global_step: 33 | reduced_train_loss: 0.5352 | train_step_timing in s: 2.129 | consumed_samples: 2176 | val_loss: 0.477\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:31 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 34/99 | lr: 3.431e-06 | global_batch_size: 64 | global_step: 34 | reduced_train_loss: 0.635 | train_step_timing in s: 0.7587 | consumed_samples: 2240 | val_loss: 0.477\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:32 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 35/99 | lr: 3.529e-06 | global_batch_size: 64 | global_step: 35 | reduced_train_loss: 0.6193 | train_step_timing in s: 0.8925 | consumed_samples: 2304 | val_loss: 0.477\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:34 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 36/99 | lr: 3.627e-06 | global_batch_size: 64 | global_step: 36 | reduced_train_loss: 0.5338 | train_step_timing in s: 0.83 | consumed_samples: 2368 | val_loss: 0.477\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:35 nemo_logging:393] Successfully saved checkpoint from iteration      30 to /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.48-step=29-consumed_samples=1920.0-last.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:35 nemo_logging:393] Async checkpoint save for step 30 (/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.48-step=29-consumed_samples=1920.0-last.ckpt) finalized successfully.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:37 nemo_logging:393] Async finalization time took 2.090 s\n",
      "i.finetune/0 Training epoch 0, iteration 37/99 | lr: 3.725e-06 | global_batch_size: 64 | global_step: 37 | reduced_train_loss: 0.5765 | train_step_timing in s: 0.6128 | consumed_samples: 2432 | val_loss: 0.477\n",
      "i.finetune/0 Training epoch 0, iteration 38/99 | lr: 3.824e-06 | global_batch_size: 64 | global_step: 38 | reduced_train_loss: 0.5206 | train_step_timing in s: 0.6144 | consumed_samples: 2496 | val_loss: 0.477\n",
      "i.finetune/0 Training epoch 0, iteration 39/99 | lr: 3.922e-06 | global_batch_size: 64 | global_step: 39 | reduced_train_loss: 0.5687 | train_step_timing in s: 0.6094 | consumed_samples: 2560 | val_loss: 0.477\n",
      "i.finetune/0 Validation: iteration 1/0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:43 nemo_logging:393] Global Checkpoint Save : Rank: 0 : Iteration: 40 : Start time: 1756196203.789s : Save duration: 0.178s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:44 nemo_logging:393] Scheduled async checkpoint save for /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.33-step=39-consumed_samples=2560.0-last.ckpt\n",
      "i.finetune/0 Training epoch 0, iteration 40/99 | lr: 4.02e-06 | global_batch_size: 64 | global_step: 40 | reduced_train_loss: 0.6109 | train_step_timing in s: 2.64 | consumed_samples: 2624 | val_loss: 0.3335\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:47 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 41/99 | lr: 4.118e-06 | global_batch_size: 64 | global_step: 41 | reduced_train_loss: 0.4909 | train_step_timing in s: 1.816 | consumed_samples: 2688 | val_loss: 0.3335\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:49 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 42/99 | lr: 4.216e-06 | global_batch_size: 64 | global_step: 42 | reduced_train_loss: 0.3827 | train_step_timing in s: 2.562 | consumed_samples: 2752 | val_loss: 0.3335\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:51 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 43/99 | lr: 4.314e-06 | global_batch_size: 64 | global_step: 43 | reduced_train_loss: 0.5915 | train_step_timing in s: 1.195 | consumed_samples: 2816 | val_loss: 0.3335\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:53 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 44/99 | lr: 4.412e-06 | global_batch_size: 64 | global_step: 44 | reduced_train_loss: 0.4703 | train_step_timing in s: 1.514 | consumed_samples: 2880 | val_loss: 0.3335\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:55 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 45/99 | lr: 4.51e-06 | global_batch_size: 64 | global_step: 45 | reduced_train_loss: 0.4229 | train_step_timing in s: 0.6167 | consumed_samples: 2944 | val_loss: 0.3335\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:56 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 46/99 | lr: 4.608e-06 | global_batch_size: 64 | global_step: 46 | reduced_train_loss: 0.2295 | train_step_timing in s: 1.527 | consumed_samples: 3008 | val_loss: 0.3335\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:58 nemo_logging:393] Successfully saved checkpoint from iteration      40 to /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.33-step=39-consumed_samples=2560.0-last.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:58 nemo_logging:393] Async checkpoint save for step 40 (/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.33-step=39-consumed_samples=2560.0-last.ckpt) finalized successfully.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:46:59 nemo_logging:393] Async finalization time took 1.631 s\n",
      "i.finetune/0 Training epoch 0, iteration 47/99 | lr: 4.706e-06 | global_batch_size: 64 | global_step: 47 | reduced_train_loss: 0.386 | train_step_timing in s: 1.227 | consumed_samples: 3072 | val_loss: 0.3335\n",
      "i.finetune/0 Training epoch 0, iteration 48/99 | lr: 4.804e-06 | global_batch_size: 64 | global_step: 48 | reduced_train_loss: 0.3955 | train_step_timing in s: 0.6098 | consumed_samples: 3136 | val_loss: 0.3335\n",
      "i.finetune/0 Training epoch 0, iteration 49/99 | lr: 4.902e-06 | global_batch_size: 64 | global_step: 49 | reduced_train_loss: 0.3621 | train_step_timing in s: 0.6074 | consumed_samples: 3200 | val_loss: 0.3335\n",
      "i.finetune/0 INFO:pytorch_lightning.utilities.rank_zero:Epoch 0, global step 49: 'val_loss' reached 0.33346 (best 0.33346), saving model to '/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.33-step=49-consumed_samples=3200.0.ckpt' as top 2\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:03 nemo_logging:393] Global Checkpoint Save : Rank: 0 : Iteration: 49 : Start time: 1756196222.392s : Save duration: 1.369s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:04 nemo_logging:393] Scheduled async checkpoint save for /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.33-step=49-consumed_samples=3200.0.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:04 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Validation: iteration 1/0\n",
      "i.finetune/0 Training epoch 0, iteration 50/99 | lr: 5e-06 | global_batch_size: 64 | global_step: 50 | reduced_train_loss: 0.3359 | train_step_timing in s: 1.452 | consumed_samples: 3264 | val_loss: 0.2905\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:14 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 51/99 | lr: 4.996e-06 | global_batch_size: 64 | global_step: 51 | reduced_train_loss: 0.4107 | train_step_timing in s: 1.338 | consumed_samples: 3328 | val_loss: 0.2905\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:16 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 52/99 | lr: 4.982e-06 | global_batch_size: 64 | global_step: 52 | reduced_train_loss: 0.4784 | train_step_timing in s: 0.6187 | consumed_samples: 3392 | val_loss: 0.2905\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:17 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 53/99 | lr: 4.96e-06 | global_batch_size: 64 | global_step: 53 | reduced_train_loss: 0.3316 | train_step_timing in s: 0.9179 | consumed_samples: 3456 | val_loss: 0.2905\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:18 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 54/99 | lr: 4.929e-06 | global_batch_size: 64 | global_step: 54 | reduced_train_loss: 0.2946 | train_step_timing in s: 0.7666 | consumed_samples: 3520 | val_loss: 0.2905\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:19 nemo_logging:393] Successfully saved checkpoint from iteration      49 to /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.33-step=49-consumed_samples=3200.0.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:19 nemo_logging:393] Async checkpoint save for step 50 (/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.33-step=49-consumed_samples=3200.0.ckpt) finalized successfully.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:21 nemo_logging:393] Async finalization time took 2.072 s\n",
      "i.finetune/0 Training epoch 0, iteration 55/99 | lr: 4.89e-06 | global_batch_size: 64 | global_step: 55 | reduced_train_loss: 0.4005 | train_step_timing in s: 0.6211 | consumed_samples: 3584 | val_loss: 0.2905\n",
      "i.finetune/0 Training epoch 0, iteration 56/99 | lr: 4.842e-06 | global_batch_size: 64 | global_step: 56 | reduced_train_loss: 0.2703 | train_step_timing in s: 0.6096 | consumed_samples: 3648 | val_loss: 0.2905\n",
      "i.finetune/0 Training epoch 0, iteration 57/99 | lr: 4.786e-06 | global_batch_size: 64 | global_step: 57 | reduced_train_loss: 0.3257 | train_step_timing in s: 1.494 | consumed_samples: 3712 | val_loss: 0.2905\n",
      "i.finetune/0 Training epoch 0, iteration 58/99 | lr: 4.722e-06 | global_batch_size: 64 | global_step: 58 | reduced_train_loss: 0.1669 | train_step_timing in s: 0.6648 | consumed_samples: 3776 | val_loss: 0.2905\n",
      "i.finetune/0 Training epoch 0, iteration 59/99 | lr: 4.65e-06 | global_batch_size: 64 | global_step: 59 | reduced_train_loss: 0.3042 | train_step_timing in s: 0.6153 | consumed_samples: 3840 | val_loss: 0.2905\n",
      "i.finetune/0 Validation: iteration 1/0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:29 nemo_logging:393] Global Checkpoint Save : Rank: 0 : Iteration: 60 : Start time: 1756196249.805s : Save duration: 0.180s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:31 nemo_logging:393] Scheduled async checkpoint save for /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.27-step=59-consumed_samples=3840.0-last.ckpt\n",
      "i.finetune/0 Training epoch 0, iteration 60/99 | lr: 4.57e-06 | global_batch_size: 64 | global_step: 60 | reduced_train_loss: 0.3881 | train_step_timing in s: 1.892 | consumed_samples: 3904 | val_loss: 0.2716\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:32 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 61/99 | lr: 4.484e-06 | global_batch_size: 64 | global_step: 61 | reduced_train_loss: 0.4002 | train_step_timing in s: 2.633 | consumed_samples: 3968 | val_loss: 0.2716\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:35 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 62/99 | lr: 4.39e-06 | global_batch_size: 64 | global_step: 62 | reduced_train_loss: 0.3023 | train_step_timing in s: 1.857 | consumed_samples: 4032 | val_loss: 0.2716\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:37 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 63/99 | lr: 4.29e-06 | global_batch_size: 64 | global_step: 63 | reduced_train_loss: 0.5401 | train_step_timing in s: 1.121 | consumed_samples: 4096 | val_loss: 0.2716\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:40 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 64/99 | lr: 4.184e-06 | global_batch_size: 64 | global_step: 64 | reduced_train_loss: 0.3145 | train_step_timing in s: 0.794 | consumed_samples: 4160 | val_loss: 0.2716\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:41 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 65/99 | lr: 4.073e-06 | global_batch_size: 64 | global_step: 65 | reduced_train_loss: 0.2081 | train_step_timing in s: 0.6096 | consumed_samples: 4224 | val_loss: 0.2716\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:42 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 66/99 | lr: 3.956e-06 | global_batch_size: 64 | global_step: 66 | reduced_train_loss: 0.3903 | train_step_timing in s: 1.128 | consumed_samples: 4288 | val_loss: 0.2716\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:44 nemo_logging:393] Successfully saved checkpoint from iteration      60 to /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.27-step=59-consumed_samples=3840.0-last.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:44 nemo_logging:393] Async checkpoint save for step 60 (/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.27-step=59-consumed_samples=3840.0-last.ckpt) finalized successfully.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:44 nemo_logging:393] Async finalization time took 0.067 s\n",
      "i.finetune/0 Training epoch 0, iteration 67/99 | lr: 3.834e-06 | global_batch_size: 64 | global_step: 67 | reduced_train_loss: 0.2715 | train_step_timing in s: 0.8455 | consumed_samples: 4352 | val_loss: 0.2716\n",
      "i.finetune/0 Training epoch 0, iteration 68/99 | lr: 3.708e-06 | global_batch_size: 64 | global_step: 68 | reduced_train_loss: 0.3322 | train_step_timing in s: 0.6099 | consumed_samples: 4416 | val_loss: 0.2716\n",
      "i.finetune/0 Training epoch 0, iteration 69/99 | lr: 3.578e-06 | global_batch_size: 64 | global_step: 69 | reduced_train_loss: 0.2033 | train_step_timing in s: 0.6547 | consumed_samples: 4480 | val_loss: 0.2716\n",
      "i.finetune/0 Validation: iteration 1/0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:50 nemo_logging:393] Global Checkpoint Save : Rank: 0 : Iteration: 70 : Start time: 1756196270.422s : Save duration: 0.183s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:51 nemo_logging:393] Scheduled async checkpoint save for /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.26-step=69-consumed_samples=4480.0-last.ckpt\n",
      "i.finetune/0 Training epoch 0, iteration 70/99 | lr: 3.445e-06 | global_batch_size: 64 | global_step: 70 | reduced_train_loss: 0.3736 | train_step_timing in s: 2.667 | consumed_samples: 4544 | val_loss: 0.2568\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:53 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 71/99 | lr: 3.31e-06 | global_batch_size: 64 | global_step: 71 | reduced_train_loss: 0.3343 | train_step_timing in s: 2.345 | consumed_samples: 4608 | val_loss: 0.2568\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:56 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 72/99 | lr: 3.172e-06 | global_batch_size: 64 | global_step: 72 | reduced_train_loss: 0.2417 | train_step_timing in s: 2.278 | consumed_samples: 4672 | val_loss: 0.2568\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:47:58 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 73/99 | lr: 3.032e-06 | global_batch_size: 64 | global_step: 73 | reduced_train_loss: 0.2845 | train_step_timing in s: 2.14 | consumed_samples: 4736 | val_loss: 0.2568\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:00 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 74/99 | lr: 2.891e-06 | global_batch_size: 64 | global_step: 74 | reduced_train_loss: 0.2358 | train_step_timing in s: 0.7182 | consumed_samples: 4800 | val_loss: 0.2568\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:01 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 75/99 | lr: 2.75e-06 | global_batch_size: 64 | global_step: 75 | reduced_train_loss: 0.2576 | train_step_timing in s: 0.9539 | consumed_samples: 4864 | val_loss: 0.2568\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:03 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 76/99 | lr: 2.609e-06 | global_batch_size: 64 | global_step: 76 | reduced_train_loss: 0.4009 | train_step_timing in s: 0.9047 | consumed_samples: 4928 | val_loss: 0.2568\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:05 nemo_logging:393] Successfully saved checkpoint from iteration      70 to /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.26-step=69-consumed_samples=4480.0-last.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:05 nemo_logging:393] Async checkpoint save for step 70 (/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.26-step=69-consumed_samples=4480.0-last.ckpt) finalized successfully.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:07 nemo_logging:393] Async finalization time took 2.102 s\n",
      "i.finetune/0 Training epoch 0, iteration 77/99 | lr: 2.468e-06 | global_batch_size: 64 | global_step: 77 | reduced_train_loss: 0.3106 | train_step_timing in s: 0.6625 | consumed_samples: 4992 | val_loss: 0.2568\n",
      "i.finetune/0 Training epoch 0, iteration 78/99 | lr: 2.328e-06 | global_batch_size: 64 | global_step: 78 | reduced_train_loss: 0.3653 | train_step_timing in s: 0.8748 | consumed_samples: 5056 | val_loss: 0.2568\n",
      "i.finetune/0 Training epoch 0, iteration 79/99 | lr: 2.19e-06 | global_batch_size: 64 | global_step: 79 | reduced_train_loss: 0.3352 | train_step_timing in s: 0.6022 | consumed_samples: 5120 | val_loss: 0.2568\n",
      "i.finetune/0 Validation: iteration 1/0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:14 nemo_logging:393] Global Checkpoint Save : Rank: 0 : Iteration: 80 : Start time: 1756196293.468s : Save duration: 0.838s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:14 nemo_logging:393] Scheduled async checkpoint save for /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.25-step=79-consumed_samples=5120.0-last.ckpt\n",
      "i.finetune/0 Training epoch 0, iteration 80/99 | lr: 2.055e-06 | global_batch_size: 64 | global_step: 80 | reduced_train_loss: 0.2359 | train_step_timing in s: 2.118 | consumed_samples: 5184 | val_loss: 0.2471\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:16 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 81/99 | lr: 1.922e-06 | global_batch_size: 64 | global_step: 81 | reduced_train_loss: 0.3469 | train_step_timing in s: 2.529 | consumed_samples: 5248 | val_loss: 0.2471\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:19 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 82/99 | lr: 1.792e-06 | global_batch_size: 64 | global_step: 82 | reduced_train_loss: 0.256 | train_step_timing in s: 1.771 | consumed_samples: 5312 | val_loss: 0.2471\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:21 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 83/99 | lr: 1.666e-06 | global_batch_size: 64 | global_step: 83 | reduced_train_loss: 0.2272 | train_step_timing in s: 2.112 | consumed_samples: 5376 | val_loss: 0.2471\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:24 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 84/99 | lr: 1.544e-06 | global_batch_size: 64 | global_step: 84 | reduced_train_loss: 0.1708 | train_step_timing in s: 0.9113 | consumed_samples: 5440 | val_loss: 0.2471\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:25 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 85/99 | lr: 1.427e-06 | global_batch_size: 64 | global_step: 85 | reduced_train_loss: 0.2699 | train_step_timing in s: 1.229 | consumed_samples: 5504 | val_loss: 0.2471\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:26 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 86/99 | lr: 1.316e-06 | global_batch_size: 64 | global_step: 86 | reduced_train_loss: 0.2698 | train_step_timing in s: 0.9413 | consumed_samples: 5568 | val_loss: 0.2471\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:28 nemo_logging:393] Successfully saved checkpoint from iteration      80 to /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.25-step=79-consumed_samples=5120.0-last.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:28 nemo_logging:393] Async checkpoint save for step 80 (/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.25-step=79-consumed_samples=5120.0-last.ckpt) finalized successfully.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:30 nemo_logging:393] Async finalization time took 2.175 s\n",
      "i.finetune/0 Training epoch 0, iteration 87/99 | lr: 1.21e-06 | global_batch_size: 64 | global_step: 87 | reduced_train_loss: 0.4019 | train_step_timing in s: 0.6013 | consumed_samples: 5632 | val_loss: 0.2471\n",
      "i.finetune/0 Training epoch 0, iteration 88/99 | lr: 1.11e-06 | global_batch_size: 64 | global_step: 88 | reduced_train_loss: 0.4198 | train_step_timing in s: 0.61 | consumed_samples: 5696 | val_loss: 0.2471\n",
      "i.finetune/0 Training epoch 0, iteration 89/99 | lr: 1.016e-06 | global_batch_size: 64 | global_step: 89 | reduced_train_loss: 0.1869 | train_step_timing in s: 0.5991 | consumed_samples: 5760 | val_loss: 0.2471\n",
      "i.finetune/0 Validation: iteration 1/0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:36 nemo_logging:393] Global Checkpoint Save : Rank: 0 : Iteration: 90 : Start time: 1756196316.487s : Save duration: 0.182s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:37 nemo_logging:393] Scheduled async checkpoint save for /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.24-step=89-consumed_samples=5760.0-last.ckpt\n",
      "i.finetune/0 Training epoch 0, iteration 90/99 | lr: 9.297e-07 | global_batch_size: 64 | global_step: 90 | reduced_train_loss: 0.3109 | train_step_timing in s: 2.622 | consumed_samples: 5824 | val_loss: 0.244\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:39 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 91/99 | lr: 8.503e-07 | global_batch_size: 64 | global_step: 91 | reduced_train_loss: 0.2357 | train_step_timing in s: 2.118 | consumed_samples: 5888 | val_loss: 0.244\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:41 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 92/99 | lr: 7.783e-07 | global_batch_size: 64 | global_step: 92 | reduced_train_loss: 0.3269 | train_step_timing in s: 2.29 | consumed_samples: 5952 | val_loss: 0.244\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:44 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 93/99 | lr: 7.141e-07 | global_batch_size: 64 | global_step: 93 | reduced_train_loss: 0.2552 | train_step_timing in s: 1.486 | consumed_samples: 6016 | val_loss: 0.244\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:46 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 94/99 | lr: 6.58e-07 | global_batch_size: 64 | global_step: 94 | reduced_train_loss: 0.14 | train_step_timing in s: 0.8018 | consumed_samples: 6080 | val_loss: 0.244\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:47 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 95/99 | lr: 6.101e-07 | global_batch_size: 64 | global_step: 95 | reduced_train_loss: 0.2713 | train_step_timing in s: 0.8777 | consumed_samples: 6144 | val_loss: 0.244\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:49 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Training epoch 0, iteration 96/99 | lr: 5.707e-07 | global_batch_size: 64 | global_step: 96 | reduced_train_loss: 0.3404 | train_step_timing in s: 0.8661 | consumed_samples: 6208 | val_loss: 0.244\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:50 nemo_logging:393] Successfully saved checkpoint from iteration      90 to /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.24-step=89-consumed_samples=5760.0-last.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:50 nemo_logging:393] Async checkpoint save for step 90 (/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.24-step=89-consumed_samples=5760.0-last.ckpt) finalized successfully.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:52 nemo_logging:393] Async finalization time took 1.898 s\n",
      "i.finetune/0 Training epoch 0, iteration 97/99 | lr: 5.399e-07 | global_batch_size: 64 | global_step: 97 | reduced_train_loss: 0.3045 | train_step_timing in s: 0.6424 | consumed_samples: 6272 | val_loss: 0.244\n",
      "i.finetune/0 Training epoch 0, iteration 98/99 | lr: 5.177e-07 | global_batch_size: 64 | global_step: 98 | reduced_train_loss: 0.4598 | train_step_timing in s: 0.6063 | consumed_samples: 6336 | val_loss: 0.244\n",
      "i.finetune/0 Training epoch 0, iteration 99/99 | lr: 5.044e-07 | global_batch_size: 64 | global_step: 99 | reduced_train_loss: 0.3124 | train_step_timing in s: 0.6147 | consumed_samples: 6400 | val_loss: 0.244\n",
      "i.finetune/0 INFO:pytorch_lightning.utilities.rank_zero:Epoch 0, global step 99: 'val_loss' reached 0.24404 (best 0.24404), saving model to '/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.24-step=99-consumed_samples=6400.0.ckpt' as top 2\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:55 nemo_logging:393] Global Checkpoint Save : Rank: 0 : Iteration: 99 : Start time: 1756196334.360s : Save duration: 0.755s\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:55 nemo_logging:393] Scheduled async checkpoint save for /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.24-step=99-consumed_samples=6400.0.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:48:55 nemo_logging:393] Async finalization time took 0.000 s\n",
      "i.finetune/0 Validation: iteration 1/0\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:49:00 nemo_logging:393] Successfully saved checkpoint from iteration      99 to /host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.24-step=99-consumed_samples=6400.0.ckpt\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:49:00 nemo_logging:393] Async checkpoint save for step 100 (/host_pwd/PythonNotebook/nemo_experiments/allnli_e5_large_finetune/2025-08-26_13-44-20/checkpoints/model_name=0--val_loss=0.24-step=99-consumed_samples=6400.0.ckpt) finalized successfully.\n",
      "i.finetune/0 [NeMo I 2025-08-26 13:49:02 nemo_logging:393] Async finalization time took 2.160 s\n",
      "i.finetune/0 INFO:pytorch_lightning.utilities.rank_zero:`Trainer.fit` stopped: `max_steps=100` reached.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Job nemo.collections.llm.api.finetune-zn4250k9j945t finished: SUCCEEDED\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"background-color: #272822\">                                                                                                                   </span>\n",
       "<span style=\"color: #959077; text-decoration-color: #959077; background-color: #272822\"># The experiment was run with the following tasks: ['nemo.collections.llm.api.finetune']</span><span style=\"background-color: #272822\">                           </span>\n",
       "<span style=\"color: #959077; text-decoration-color: #959077; background-color: #272822\"># You can inspect and reconstruct this experiment at a later point in time using:</span><span style=\"background-color: #272822\">                                  </span>\n",
       "<span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">experiment </span><span style=\"color: #ff4689; text-decoration-color: #ff4689; background-color: #272822\">=</span><span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\"> run</span><span style=\"color: #ff4689; text-decoration-color: #ff4689; background-color: #272822\">.</span><span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">Experiment</span><span style=\"color: #ff4689; text-decoration-color: #ff4689; background-color: #272822\">.</span><span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">from_id(</span><span style=\"color: #e6db74; text-decoration-color: #e6db74; background-color: #272822\">\"nemo.collections.llm.api.finetune_1756196008\"</span><span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">)</span><span style=\"background-color: #272822\">                                </span>\n",
       "<span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">experiment</span><span style=\"color: #ff4689; text-decoration-color: #ff4689; background-color: #272822\">.</span><span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">status() </span><span style=\"color: #959077; text-decoration-color: #959077; background-color: #272822\"># Gets the overall status</span><span style=\"background-color: #272822\">                                                                      </span>\n",
       "<span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">experiment</span><span style=\"color: #ff4689; text-decoration-color: #ff4689; background-color: #272822\">.</span><span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">logs(</span><span style=\"color: #e6db74; text-decoration-color: #e6db74; background-color: #272822\">\"nemo.collections.llm.api.finetune\"</span><span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">) </span><span style=\"color: #959077; text-decoration-color: #959077; background-color: #272822\"># Gets the log for the provided task</span><span style=\"background-color: #272822\">                          </span>\n",
       "<span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">experiment</span><span style=\"color: #ff4689; text-decoration-color: #ff4689; background-color: #272822\">.</span><span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">cancel(</span><span style=\"color: #e6db74; text-decoration-color: #e6db74; background-color: #272822\">\"nemo.collections.llm.api.finetune\"</span><span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">) </span><span style=\"color: #959077; text-decoration-color: #959077; background-color: #272822\"># Cancels the provided task if still running</span><span style=\"background-color: #272822\">                </span>\n",
       "<span style=\"background-color: #272822\">                                                                                                                   </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[48;2;39;40;34m                                                                                                                   \u001b[0m\n",
       "\u001b[38;2;149;144;119;48;2;39;40;34m# The experiment was run with the following tasks: ['nemo.collections.llm.api.finetune']\u001b[0m\u001b[48;2;39;40;34m                           \u001b[0m\n",
       "\u001b[38;2;149;144;119;48;2;39;40;34m# You can inspect and reconstruct this experiment at a later point in time using:\u001b[0m\u001b[48;2;39;40;34m                                  \u001b[0m\n",
       "\u001b[38;2;248;248;242;48;2;39;40;34mexperiment\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;255;70;137;48;2;39;40;34m=\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mrun\u001b[0m\u001b[38;2;255;70;137;48;2;39;40;34m.\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mExperiment\u001b[0m\u001b[38;2;255;70;137;48;2;39;40;34m.\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mfrom_id\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m(\u001b[0m\u001b[38;2;230;219;116;48;2;39;40;34m\"\u001b[0m\u001b[38;2;230;219;116;48;2;39;40;34mnemo.collections.llm.api.finetune_1756196008\u001b[0m\u001b[38;2;230;219;116;48;2;39;40;34m\"\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m)\u001b[0m\u001b[48;2;39;40;34m                                \u001b[0m\n",
       "\u001b[38;2;248;248;242;48;2;39;40;34mexperiment\u001b[0m\u001b[38;2;255;70;137;48;2;39;40;34m.\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mstatus\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m(\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m)\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;149;144;119;48;2;39;40;34m# Gets the overall status\u001b[0m\u001b[48;2;39;40;34m                                                                      \u001b[0m\n",
       "\u001b[38;2;248;248;242;48;2;39;40;34mexperiment\u001b[0m\u001b[38;2;255;70;137;48;2;39;40;34m.\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mlogs\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m(\u001b[0m\u001b[38;2;230;219;116;48;2;39;40;34m\"\u001b[0m\u001b[38;2;230;219;116;48;2;39;40;34mnemo.collections.llm.api.finetune\u001b[0m\u001b[38;2;230;219;116;48;2;39;40;34m\"\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m)\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;149;144;119;48;2;39;40;34m# Gets the log for the provided task\u001b[0m\u001b[48;2;39;40;34m                          \u001b[0m\n",
       "\u001b[38;2;248;248;242;48;2;39;40;34mexperiment\u001b[0m\u001b[38;2;255;70;137;48;2;39;40;34m.\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mcancel\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m(\u001b[0m\u001b[38;2;230;219;116;48;2;39;40;34m\"\u001b[0m\u001b[38;2;230;219;116;48;2;39;40;34mnemo.collections.llm.api.finetune\u001b[0m\u001b[38;2;230;219;116;48;2;39;40;34m\"\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m)\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;149;144;119;48;2;39;40;34m# Cancels the provided task if still running\u001b[0m\u001b[48;2;39;40;34m                \u001b[0m\n",
       "\u001b[48;2;39;40;34m                                                                                                                   \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"background-color: #272822\">                                                                                                                   </span>\n",
       "<span style=\"color: #959077; text-decoration-color: #959077; background-color: #272822\"># You can inspect this experiment at a later point in time using the CLI as well:</span><span style=\"background-color: #272822\">                                  </span>\n",
       "<span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">nemo experiment status nemo.collections.llm.api.finetune_1756196008</span><span style=\"background-color: #272822\">                                                </span>\n",
       "<span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">nemo experiment logs nemo.collections.llm.api.finetune_1756196008 </span><span style=\"color: #ae81ff; text-decoration-color: #ae81ff; background-color: #272822\">0</span><span style=\"background-color: #272822\">                                                </span>\n",
       "<span style=\"color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822\">nemo experiment cancel nemo.collections.llm.api.finetune_1756196008 </span><span style=\"color: #ae81ff; text-decoration-color: #ae81ff; background-color: #272822\">0</span><span style=\"background-color: #272822\">                                              </span>\n",
       "<span style=\"background-color: #272822\">                                                                                                                   </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[48;2;39;40;34m                                                                                                                   \u001b[0m\n",
       "\u001b[38;2;149;144;119;48;2;39;40;34m# You can inspect this experiment at a later point in time using the CLI as well:\u001b[0m\u001b[48;2;39;40;34m                                  \u001b[0m\n",
       "\u001b[38;2;248;248;242;48;2;39;40;34mnemo\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mexperiment\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mstatus\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mnemo.collections.llm.api.finetune_1756196008\u001b[0m\u001b[48;2;39;40;34m                                                \u001b[0m\n",
       "\u001b[38;2;248;248;242;48;2;39;40;34mnemo\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mexperiment\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mlogs\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mnemo.collections.llm.api.finetune_1756196008\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;174;129;255;48;2;39;40;34m0\u001b[0m\u001b[48;2;39;40;34m                                                \u001b[0m\n",
       "\u001b[38;2;248;248;242;48;2;39;40;34mnemo\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mexperiment\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mcancel\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34mnemo.collections.llm.api.finetune_1756196008\u001b[0m\u001b[38;2;248;248;242;48;2;39;40;34m \u001b[0m\u001b[38;2;174;129;255;48;2;39;40;34m0\u001b[0m\u001b[48;2;39;40;34m                                              \u001b[0m\n",
       "\u001b[48;2;39;40;34m                                                                                                                   \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_e5_on_allnli(TRAIN_DATA_PATH)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
